#pragma once

#include <atomic>
#include <functional>
#include <map>
#include <memory>

#include "google/protobuf/service.h"
#include "netlib/base/mutex.h"
#include "netlib/rpc/protobuf_codec.h"

namespace netlib::rpc {

class RpcMessage;
using RpcCodec = ProtobufCodec<RpcMessage>;
using ProtobufMessageCallback = std::function<void(
    const net::TcpConnectionPtr& conn, std::unique_ptr<ProtobufMessage>&& msg, Timestamp ts)>;

class RpcChannelImpl : public ::google::protobuf::RpcChannel {
public:
	explicit RpcChannelImpl(net::TcpConnectionPtr conn = nullptr) :
	    codec_(std::make_unique<RpcCodec>(
	        [this](const net::TcpConnectionPtr& conn,
	               std::unique_ptr<ProtobufMessage>&& msg,
	               Timestamp ts) { OnProtobufMessage(conn, std::move(msg), ts); })),
	    conn_(std::move(conn)) {}

	void CallMethod(const ::google::protobuf::MethodDescriptor* method,
	                ::google::protobuf::RpcController* controller,
	                const ProtobufMessage* request,
	                ProtobufMessage* response,
	                ::google::protobuf::Closure* done) override;

	void SetConnection(const net::TcpConnectionPtr& conn) { conn_ = conn; }
	void OnMessage(const net::TcpConnectionPtr& conn, net::Buffer* buf, Timestamp ts) {
		codec_->OnMessage(conn, buf, ts);
	}

private:
	void OnProtobufMessage(const net::TcpConnectionPtr& conn,
	                       std::unique_ptr<ProtobufMessage>&& msg,
	                       Timestamp ts);

private:
	std::atomic<int64_t> id_creator_{0};
	std::unique_ptr<RpcCodec> codec_;
	net::TcpConnectionPtr conn_; // NOLINT
	using RequestResult = std::pair<ProtobufMessage*, ::google::protobuf::Closure*>;
	std::map<int64_t, RequestResult> req_res_map_;
	mutable MutexLock mtx_;
};

} // namespace netlib::rpc